import requests
from flask import Blueprint, render_template, request
from .ygy_data_out import get_word_cloud_data_list
from .ygy_sd_api import text2img_base64 as api_text2img_base64
from .ygy_sd_api import set_model as api_set_model
from .ygy_sd_api.api_test import model_name_list

ygy = Blueprint('ygy', __name__, template_folder='ygy_templates', static_folder='ygy_static')
res_dict = {
    'data': {},
    'msg': '',
    'status': 0
}


def get_error_dict(msg: str):
    result_dict = res_dict
    result_dict['msg'] = msg
    return result_dict


def get_success_dict(data: dict, msg: str) -> dict:
    result_dict = res_dict
    result_dict['data'] = data
    result_dict['msg'] = msg
    result_dict['status'] = 1
    return result_dict


@ygy.route('/')
def index_page():
    model_list = enumerate(model_name_list)
    num_list = range(1, len(model_name_list) + 1)
    size_list = ["512*512", "512*768", "768*512", "768*768"]
    return render_template('ygy_sd_ai.html', model_list=model_list, num_list=num_list, size_list=size_list)


@ygy.route('/page1')
def page1():
    return render_template('ygy_page1.html')


@ygy.route('/data_out/tags_word_cloud/<int:top_num>', methods=['GET'])
def tags_word_cloud(top_num: int):
    try:
        res_list = get_word_cloud_data_list(top_num)
    except Exception as e:
        print(e)
        return get_error_dict('未知错误！')
    else:
        return get_success_dict({"list": res_list}, '获取成功！')


@ygy.route('/sd_api/text2img', methods=['POST'])
def text2img():
    post_data = request.get_json()
    try:
        img_base64_list = api_text2img_base64(
            post_data['prompt'], post_data['gen_num'],
            post_data['gen_width'], post_data['gen_height'])
    except requests.exceptions.RequestException:
        return get_error_dict('服务器链接错误！')
    except KeyError:
        return get_error_dict('参数非法！')
    except ValueError:
        return get_error_dict('参数取值非法！')
    else:
        return get_success_dict({'list': img_base64_list}, '图像生成成功')


@ygy.route('/sd_api/set_model', methods=['POST'])
def set_model():
    post_data = request.get_json()
    try:
        api_set_model(post_data['model_idx'])
    except requests.exceptions.RequestException:
        return get_error_dict('服务器链接错误！')
    except KeyError:
        return get_error_dict('参数非法！')
    except ValueError:
        return get_error_dict('参数取值非法！')
    else:
        return get_success_dict({}, '模型修改成功')
